#include "crypto_common.h"

// 逐步测试每个加密步骤的可逆性
void test_each_step() {
    crypto_context_t ctx;
    unsigned char test_data[16] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 
                                   0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F};
    unsigned char backup[16];
    unsigned char step_data[16];
    
    printf("=== 逐步测试每个加密步骤 ===\n");
    
    // 初始化上下文
    secure_random_init();
    ctx.complexity_level = 3;
    ctx.seed = 12345;
    
    for (int i = 0; i < SALT_LENGTH; i++) {
        ctx.salt[i] = (unsigned char)(i % 256);
    }
    
    for (int i = 0; i < KEY_LENGTH; i++) {
        ctx.primary_key[i] = (unsigned char)((i * 37) % 256);
    }
    
    derive_keys(&ctx);
    
    memcpy(backup, test_data, 16);
    
    printf("原始数据: ");
    for (int i = 0; i < 16; i++) printf("%02x ", test_data[i]);
    printf("\n\n");
    
    // 测试单轮加密的每一步
    int round = 0;
    memcpy(step_data, test_data, 16);
    
    printf("=== 测试第 %d 轮加密的每一步 ===\n", round);
    
    // 第1步：XOR加密
    printf("第1步: XOR加密\n");
    printf("加密前: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    xor_block(step_data, ctx.derived_keys[round], 16);
    
    printf("加密后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    // 逆向第1步
    xor_block(step_data, ctx.derived_keys[round], 16);
    
    printf("逆向后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf(" -> ");
    
    if (memcmp(step_data, test_data, 16) == 0) {
        printf("✓ 可逆\n\n");
    } else {
        printf("✗ 不可逆\n\n");
        return;
    }
    
    // 第2步：字节替换
    printf("第2步: 字节替换\n");
    printf("替换前: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    substitute_bytes(step_data, 16, SBOX);
    
    printf("替换后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    // 逆向第2步
    inverse_substitute_bytes(step_data, 16, INV_SBOX);
    
    printf("逆向后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf(" -> ");
    
    if (memcmp(step_data, test_data, 16) == 0) {
        printf("✓ 可逆\n\n");
    } else {
        printf("✗ 不可逆\n\n");
        return;
    }
    
    // 第3步：循环移位
    printf("第3步: 循环移位 (左移 %d 位)\n", (round + 1) * 3);
    printf("移位前: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    rotate_left(step_data, 16, (round + 1) * 3);
    
    printf("移位后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    // 逆向第3步
    rotate_right(step_data, 16, (round + 1) * 3);
    
    printf("逆向后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf(" -> ");
    
    if (memcmp(step_data, test_data, 16) == 0) {
        printf("✓ 可逆\n\n");
    } else {
        printf("✗ 不可逆 - 循环移位有问题！\n\n");
        return;
    }
    
    // 第4步：再次XOR
    printf("第4步: 再次XOR (不同密钥片段)\n");
    printf("XOR前: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    for (size_t i = 0; i < 16; i++) {
        step_data[i] ^= ctx.derived_keys[round][(i + round * 17) % KEY_LENGTH];
    }
    
    printf("XOR后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    // 逆向第4步
    for (size_t i = 0; i < 16; i++) {
        step_data[i] ^= ctx.derived_keys[round][(i + round * 17) % KEY_LENGTH];
    }
    
    printf("逆向后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf(" -> ");
    
    if (memcmp(step_data, test_data, 16) == 0) {
        printf("✓ 可逆\n\n");
    } else {
        printf("✗ 不可逆\n\n");
        return;
    }
    
    // 第5步：复杂变换
    printf("第5步: 复杂变换\n");
    printf("变换前: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    for (size_t i = 0; i < 16; i++) {
        unsigned char temp = step_data[i];
        step_data[i] = ((temp << 3) | (temp >> 5)) ^ (unsigned char)(round * 19 + i * 23);
    }
    
    printf("变换后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf("\n");
    
    // 逆向第5步
    for (size_t i = 0; i < 16; i++) {
        step_data[i] ^= (unsigned char)(round * 19 + i * 23);
        unsigned char temp = step_data[i];
        step_data[i] = ((temp >> 3) | (temp << 5));
    }
    
    printf("逆向后: ");
    for (int i = 0; i < 16; i++) printf("%02x ", step_data[i]);
    printf(" -> ");
    
    if (memcmp(step_data, test_data, 16) == 0) {
        printf("✓ 可逆\n\n");
    } else {
        printf("✗ 不可逆\n\n");
        return;
    }
    
    printf("所有单步操作都是可逆的！\n");
}

// 测试循环移位函数的正确性
void test_rotate_functions() {
    printf("\n=== 测试循环移位函数 ===\n");
    
    unsigned char test_data[4] = {0x12, 0x34, 0x56, 0x78};
    unsigned char backup[4];
    
    memcpy(backup, test_data, 4);
    
    printf("原始数据: ");
    for (int i = 0; i < 4; i++) printf("%02x ", test_data[i]);
    printf("\n");
    
    // 测试不同的移位量
    for (int shift = 1; shift <= 8; shift++) {
        memcpy(test_data, backup, 4);
        
        printf("测试移位 %d 位: ", shift);
        
        // 左移再右移
        rotate_left(test_data, 4, shift);
        rotate_right(test_data, 4, shift);
        
        printf("结果: ");
        for (int i = 0; i < 4; i++) printf("%02x ", test_data[i]);
        
        if (memcmp(test_data, backup, 4) == 0) {
            printf(" ✓\n");
        } else {
            printf(" ✗\n");
        }
    }
}

int main() {
    test_rotate_functions();
    test_each_step();
    return 0;
}